/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "TemporalAA.h"
#include "DXHelper.h"
#include "DXUTCamera.h"

TemporalAA g_temporalAA;

/**
*/
TemporalAA::TemporalAA()
{
    // Scene texture data
    for( int nTexture=0; nTexture<TEXTURE_NUM; ++nTexture )
    {
        m_pSceneTexture[ nTexture ] = NULL;
        m_pSceneTextureSRV[ nTexture ] = NULL;
        m_pSceneTextureRTV[ nTexture ] = NULL;
    }

    m_pDepthTexture = NULL;
    m_pDSV = NULL;

    // MSAA Scene texture data
    for( int nMSAA=0; nMSAA<MSAA_NUM; ++nMSAA )
    {
        m_pSceneTexture_MS[ nMSAA ] = NULL;
        m_pSceneTextureRTV_MS[ nMSAA ] = NULL;
        m_pDepthTexture_MS[ nMSAA ] = NULL;
        m_pDSV_MS[ nMSAA ] = NULL;
    }

    m_nFrame = 0;

    m_mode = AA_MODE_TEMPORAL_2X_NO_BLEND;
}

/**
*/
TemporalAA::~TemporalAA()
{
}

/**
*/
HRESULT TemporalAA::OnCreateDevice( ID3D11Device* pd3dDevice, const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc )
{
    HRESULT hr;

    V_RETURN( m_temporalBlendPass.OnCreateDevice( pd3dDevice ) );

    return S_OK;
}

/**
*/
HRESULT TemporalAA::OnResizedSwapChain( ID3D11Device* pd3dDevice, IDXGISwapChain* pSwapChain, const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc )
{
    HRESULT hr = S_OK;

    // non-MSAA surfaces
    {
        DXGI_SAMPLE_DESC sampleDesc;
        sampleDesc.Count = 1;
        sampleDesc.Quality = 0;

        for( int nTexture=0; nTexture<TEXTURE_NUM; ++nTexture )
        {
            V_RETURN( CreateSurface( pd3dDevice, &m_pSceneTexture[ nTexture ], &m_pSceneTextureSRV[ nTexture ], &m_pSceneTextureRTV[ nTexture ],
                DXGI_FORMAT_R8G8B8A8_UNORM, pBackBufferSurfaceDesc->Width, pBackBufferSurfaceDesc->Height, sampleDesc ) );
        }

        V_RETURN( CreateDepthStencil( pd3dDevice, &m_pDepthTexture, &m_pDSV, DXGI_FORMAT_D24_UNORM_S8_UINT, 
            pBackBufferSurfaceDesc->Width, pBackBufferSurfaceDesc->Height, sampleDesc ) );

        m_pSceneTexture[ 0 ]->GetDesc( &m_sceneTextureTexDesc );
    }

    // MSAA-enabled surfaces
    {
        int nSampleCount = 2;
        for( int nMSAA=0; nMSAA<MSAA_NUM; ++nMSAA )
        {
            DXGI_SAMPLE_DESC sampleDesc;
            sampleDesc.Count = nSampleCount;
            sampleDesc.Quality = 0;

            V_RETURN( CreateSurface( pd3dDevice, &m_pSceneTexture_MS[ nMSAA ], NULL, &m_pSceneTextureRTV_MS[ nMSAA ],
                    DXGI_FORMAT_R8G8B8A8_UNORM, pBackBufferSurfaceDesc->Width, pBackBufferSurfaceDesc->Height, sampleDesc ) );

            V_RETURN( CreateDepthStencil( pd3dDevice, &m_pDepthTexture_MS[ nMSAA ], &m_pDSV_MS[ nMSAA ], DXGI_FORMAT_D24_UNORM_S8_UINT, 
                pBackBufferSurfaceDesc->Width, pBackBufferSurfaceDesc->Height, sampleDesc ) );

            nSampleCount *= 2;
        }
    }

    // DXUT's default Rasterizer state is MultisampleEnable=TRUE

    return hr;
}

/**
*/
void TemporalAA::OnReleasingSwapChain()
{
    // Scene texture data
    for( int nTexture=0; nTexture<TEXTURE_NUM; ++nTexture )
    {
        SAFE_RELEASE( m_pSceneTexture[ nTexture ] );
        SAFE_RELEASE( m_pSceneTextureSRV[ nTexture ] );
        SAFE_RELEASE( m_pSceneTextureRTV[ nTexture ] );
    }
    SAFE_RELEASE( m_pDepthTexture );
    SAFE_RELEASE( m_pDSV );

    // MSAA Scene texture data
    for( int nMSAA=0; nMSAA<MSAA_NUM; ++nMSAA )
    {
        SAFE_RELEASE( m_pSceneTexture_MS[ nMSAA ] );
        SAFE_RELEASE( m_pSceneTextureRTV_MS[ nMSAA ] );
        SAFE_RELEASE( m_pDepthTexture_MS[ nMSAA ] );
        SAFE_RELEASE( m_pDSV_MS[ nMSAA ] );
    }
}

/**
*/
void TemporalAA::OnDestroyDevice()
{
    m_temporalBlendPass.OnDestroyDevice();
}

/**
*/
void TemporalAA::BeginRenderScene( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext, const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc )
{
    // Save the old viewport
    UINT nViewPorts = 1;
    pd3dImmediateContext->RSGetViewports( &nViewPorts, m_vpOld );

    // Setup the viewport to match the backbuffer
    D3D11_VIEWPORT vp;
    vp.Width = static_cast< float >( pBackBufferSurfaceDesc->Width );
    vp.Height = static_cast< float >( pBackBufferSurfaceDesc->Height );
    vp.MinDepth = 0.0f;
    vp.MaxDepth = 1.0f;
    vp.TopLeftX = 0;
    vp.TopLeftY = 0;
    pd3dImmediateContext->RSSetViewports( 1, &vp );

    float ClearColor[4] = { 0.176f, 0.196f, 0.667f, 0.0f };
    //float SceneDepthClearColor[ 1 ] = { 1.0f };

    ID3D11RenderTargetView* pRTV = NULL;
    ID3D11DepthStencilView* pDSV = NULL;

    if( m_mode == AA_MODE_MSAA_2X ) {
        pRTV = m_pSceneTextureRTV_MS[ MSAA_2X ];
        pDSV = m_pDSV_MS[ MSAA_2X ];
    }
    else if( m_mode == AA_MODE_MSAA_4X ) {
        pRTV = m_pSceneTextureRTV_MS[ MSAA_4X ];
        pDSV = m_pDSV_MS[ MSAA_4X ];
    }
    else {
        pRTV = m_pSceneTextureRTV[ m_nFrame ];
        pDSV = m_pDSV;
    }

    pd3dImmediateContext->ClearRenderTargetView( pRTV, ClearColor );
    pd3dImmediateContext->ClearDepthStencilView( pDSV, D3D11_CLEAR_DEPTH, 1.0, 0x00 );

    // MRT setup
    ID3D11RenderTargetView* pRTVs[1];
    pRTVs[0] = pRTV;

    // bind to our scene texture and the main depth buffer 
    pd3dImmediateContext->OMSetRenderTargets( 1, (ID3D11RenderTargetView*const*)&pRTVs, pDSV );
}

/**
*/
void TemporalAA::EndRenderScene( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext )
{
    // Restore the old viewport
    UINT nViewPorts = 1;
    pd3dImmediateContext->RSSetViewports( nViewPorts, m_vpOld );

    ID3D11RenderTargetView* pRTVs[4];
    pRTVs[0] = NULL;
    pRTVs[1] = NULL;
    pRTVs[2] = NULL;
    pRTVs[3] = NULL;

    // bind to our scene texture and the main depth buffer 
    pd3dImmediateContext->OMSetRenderTargets( 4, (ID3D11RenderTargetView*const*)&pRTVs, NULL );
}

/**
*/
void TemporalAA::BeginTemporalBlendPass( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext, const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc )
{
    float ClearColor[4] = { 0.176f, 0.196f, 0.667f, 0.0f };
    ID3D11RenderTargetView* pRTV = DXUTGetD3D11RenderTargetView();

    ID3D11DepthStencilView* pDSV = DXUTGetD3D11DepthStencilView();
    pd3dImmediateContext->ClearDepthStencilView( pDSV, D3D11_CLEAR_DEPTH, 1.0, 0x00 );

    pd3dImmediateContext->ClearRenderTargetView( pRTV, ClearColor );
    pd3dImmediateContext->OMSetRenderTargets( 1, &pRTV, pDSV );
}

/**
*/
void TemporalAA::DoTemporalBlendPass( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext, const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc )
{
    if( m_mode == AA_MODE_NONE )
    {
        m_temporalBlendPass.SetpSceneTextureCurSRV( m_pSceneTextureSRV[ m_nFrame ] );
        m_temporalBlendPass.SetpSceneTexturePrevSRV( m_pSceneTextureSRV[ m_nFrame ] );
    }
    else if( m_mode == AA_MODE_MSAA_2X )
    {
        // resolve the MSAA surface.
        pd3dImmediateContext->ResolveSubresource( m_pSceneTexture[ m_nFrame ], 0, m_pSceneTexture_MS[ MSAA_2X ], 0, m_sceneTextureTexDesc.Format );
        m_temporalBlendPass.SetpSceneTextureCurSRV( m_pSceneTextureSRV[ m_nFrame ] );
        m_temporalBlendPass.SetpSceneTexturePrevSRV( m_pSceneTextureSRV[ m_nFrame ] );
    }
    else if( m_mode == AA_MODE_MSAA_4X )
    {
        // resolve the MSAA surface.
        pd3dImmediateContext->ResolveSubresource( m_pSceneTexture[ m_nFrame ], 0, m_pSceneTexture_MS[ MSAA_4X ], 0, m_sceneTextureTexDesc.Format );
        m_temporalBlendPass.SetpSceneTextureCurSRV( m_pSceneTextureSRV[ m_nFrame ] );
        m_temporalBlendPass.SetpSceneTexturePrevSRV( m_pSceneTextureSRV[ m_nFrame ] );
    }
    else if( m_mode == AA_MODE_TEMPORAL_2X_BLEND ) 
    {
        m_temporalBlendPass.SetpSceneTextureCurSRV( m_pSceneTextureSRV[ m_nFrame ] );
        m_temporalBlendPass.SetpSceneTexturePrevSRV( m_pSceneTextureSRV[ (m_nFrame+1)%TEXTURE_NUM ] );
    }
    else if( m_mode == AA_MODE_TEMPORAL_2X_NO_BLEND ) 
    {
        m_temporalBlendPass.SetpSceneTextureCurSRV( m_pSceneTextureSRV[ m_nFrame ] );
        m_temporalBlendPass.SetpSceneTexturePrevSRV( m_pSceneTextureSRV[ m_nFrame ] );
    }

    m_temporalBlendPass.OnFrameRender( pd3dDevice, pd3dImmediateContext, pBackBufferSurfaceDesc, 
        pBackBufferSurfaceDesc->Width, pBackBufferSurfaceDesc->Height );    
}

/**
*/
void TemporalAA::EndTemporalBlendPass( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext )
{
}

/**
    Based on the OpenGL Red Book Chapter 10.
    http://glprogramming.com/red/chapter10.html
*/
D3DXMATRIX TemporalAA::JitteredFrustum(float left, float right, float bottom,
    float top, float fNear, float fFar, float pixdx, 
    float pixdy, float eyedx, float eyedy, float focus, const CameraInfo& cameraInfo) const
{
    float xwsize, ywsize; 
    float dx, dy;

    xwsize = right - left;
    ywsize = top - bottom;
    // translate the screen space jitter distances into near clipping plane distances
    dx = -(pixdx*xwsize/cameraInfo.fBackBufferW + 
            eyedx*fNear/focus);
    dy = -(pixdy*ywsize/cameraInfo.fBackBufferH + 
            eyedy*fNear/focus);

    D3DXMATRIX mPerspective;

    D3DXMatrixPerspectiveOffCenterLH( &mPerspective, left + dx, right + dx, bottom + dy, top + dy, fNear, fFar );

    return mPerspective;
}

/**
    Based on the OpenGL Red Book Chapter 10.
    http://glprogramming.com/red/chapter10.html
    http://www.opengl.org/resources/faq/technical/transformations.htm#tran0085 (OpenGL FAQ and Troubleshooting Guide Chapter 9)
    Create glFrustum parameters from glPerspective parameters.
*/
D3DXMATRIX TemporalAA::JitteredPerspective( float fovy, float aspect, 
    float fNear, float fFar, float pixdx, float pixdy, 
    float eyedx, float eyedy, float focus, const CameraInfo& cameraInfo ) const
{
    float fov2,left,right,bottom,top;
    fov2 = fovy * 0.5f;

    top = fNear * tanf( fov2 );
    bottom = -top;
    right = top * aspect;
    left = -right;

    return JitteredFrustum (left, right, bottom, top, fNear, fFar,
        pixdx, pixdy, eyedx, eyedy, focus, cameraInfo);
}

/**
    @param cameraInfo   [in]    camera information.
    @return A subpixel jittered projection matrix.
*/
D3DXMATRIX TemporalAA::GetProjectionJitterMatrix( const CameraInfo& cameraInfo ) const
{
    float jitterNone[ 2 ] = {0.0f, 0.0f};
    // jitter amount is in screen space.
    float jitter0[ 2 ] = {0.25f, -0.25f};
    float jitter1[ 2 ] = {-0.25f, 0.25f};
    float jitter2[ 2 ] = {0.15f, -0.15f};
    float jitter3[ 2 ] = {-0.15f, 0.15f};

    float* jitter = NULL;


    if( m_mode == AA_MODE_NONE || m_mode == AA_MODE_MSAA_2X || m_mode == AA_MODE_MSAA_4X )
    {
        jitter = &jitterNone[ 0 ];
    }
    else if( m_mode == AA_MODE_TEMPORAL_2X_BLEND )
    {
        if( m_nFrame == 0 ) {
            jitter = &jitter0[ 0 ];
        }
        else {
            jitter = &jitter1[ 0 ];
        }
    }
    else if(  m_mode == AA_MODE_TEMPORAL_2X_NO_BLEND ) 
    {
        if( m_nFrame == 0 ) {
            jitter = &jitter2[ 0 ];
        }
        else {
            jitter = &jitter3[ 0 ];
        }
    }

    // eyedx, eyedy, focus end up multiplying to 0 and not affecting the end result.
    return JitteredPerspective( cameraInfo.fFov, cameraInfo.fAspect, cameraInfo.fNear, cameraInfo.fFar, jitter[ 0 ], jitter[ 1 ],
        0.0f, 0.0f, 1.0f, cameraInfo );
}

